//
// CollinsFragmentation.cxx
//
// Created by TB on 11/30/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <cstdlib> // For abs(int)
#include <iostream>
#include <sstream>

#include <boost/foreach.hpp>

#include <TCanvas.h>
#include <TH1D.h>
#include <TH2D.h>
#include <TMath.h>
#include <TRandom.h>
#include <TRotation.h>

#include "tppmc/CollinsFragmentation.h"
#include "tppmc/Enumerations.h"
#include "tppmc/KkpParton.h"
#include "tppmc/TF2J.h"

extern "C" {
   
   // Access to FORTRAN fDSS subroutine.
   void fdss_(int* hadron, int* charge, int* order, double* z,
              double* scale, int* parton, double* result);
   
   // Access to FORTRAN FRAGINI COMMON block.
   extern struct {
      double fini; 
   } fragini_;
   
} // extern "C"

namespace {

   double collins_z_dependence(double z, double a, double b, double n){
      return
      n *
      pow(z, a) *
      pow((1. - z), b) *
      pow((a + b), (a + b)) /
      (pow(a, a) * pow(b, b));
   }

   // Look for better alternative than global pointer
   PARAMETERS* gParams;

   /**
    2D function for DSS + Collins fragmentation.
    Wrapper around Fortran DSS subroutine (see fDSS.f) compatible with ROOT TF1.
    x[0] is the z value at which to evaluate the function.
    x[1] is the phi value.
    Depends on eight parameters, par[0] to par[7]:
    0: Average quark pT^2
    1: Q at which to evaluate
    2: Parton type, following KKP convention (see fDSS.f)
    3: LO or NLO function
    4: Hadron species (pion, kaon, proton or charged hadrons)
    5: Hadron charge (+, - or average of +/-)
    6: Depolarisation factor
    7: Transverse parton spin
    8: Collins parameterisation up quark M^2
    9: Collins parameterisation down quark M^2
    */
   double DSSFragFunctionCollins(Double_t* x, Double_t* par) {//const {
      // This FF is only valid for PIPLUS, PIMINUS or PIZERO
      // It should really only be used for u,ubar,d,dbar,s,sbar quarks
      // (otherwise use DSSFragFunction).
      int ih = static_cast<int>(par[4]);
      int iset = 0;            // parametrization is LO ONLY
      int charge = static_cast<int>(par[5]);    // parametrization is PIPLUS, PIMINUS,
                                                //PIZERO selected below...
      double z = x[0];
      double phi = x[1];
      double qs2 = par[1]*par[1];
      double dh;
      int qtyp = static_cast<int>(par[2]);
      double depol_factor = par[6];
      bool atPosivitityBound = static_cast<bool>(par[7]);
      double collins = 0.0; 
      static bool first = true;
      if(first) {
         fragini_.fini = 0;
         first = false; 
      } // if
      // Check to make sure we weren't called in some strange way. 
      if(par[4] not_eq 1.) {
         cerr <<
         "WARNING - DSSFragFunctionCollins called with par[4] = " <<
         par[4] << " (not PION, PION assumed)" << endl;
         ih = 1;
      } // if
      
      if((qtyp not_eq UQUARK) and
         (qtyp not_eq UBAR) and
         (qtyp not_eq DQUARK) and
         (qtyp not_eq DBAR) and
         (qtyp not_eq SQUARK) and
         (qtyp not_eq SBAR)) {
         qtyp = SQUARK; 
      } // if
      
      // Favored/Unfavored swaps: 
      // These lines enforces the favored/unfavored 
      // fragmentation equalities (see hep-ph0701006v3 eq.45,46):
      
      // u -> pi+   : FAVORED
      // dbar-> pi+ : FAVORED
      // d -> pi-   : FAVORED
      // ubar-> pi- : FAVORED
      
      // u -> pi0   : FAVORED
      // dbar-> pi0 : FAVORED
      // d -> pi0   : FAVORED
      // ubar-> pi0 : FAVORED
      
      // d -> pi+    : UNFAVORED
      // u -> pi-    : UNFAVORED
      // ubar -> pi+ : UNFAVORED
      // dbar -> pi- : UNFAVORED
      
      // s ->pi+/-/0     : UNFAVORED
      // sbar -> pi+/-/0 : UNFAVORED
      
      
      // The fragmentation functions are parametrized in terms of the fragmentation
      // function for a positive pion, and we choose the favored or unfavored fragmentation 
      // function from this by setting the quark type. 
      
      //      enum { FAVOURED, UNFAVOURED };
      
      int fqtyp = qtyp; 
      
      if(charge==1){
         // positive pion
         if(qtyp==UQUARK) fqtyp = UQUARK; 
         if(qtyp==DQUARK) fqtyp = DQUARK;
         if(qtyp==UBAR) fqtyp = DQUARK;
         if(qtyp==DBAR) fqtyp = UQUARK; 
         if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
      }
      else if(charge==-1){
         // negative pion
         if(qtyp==DQUARK) fqtyp = UQUARK;
         if(qtyp==UQUARK) fqtyp = DQUARK; 
         if(qtyp==UBAR) fqtyp = UQUARK;
         if(qtyp==DBAR) fqtyp = DQUARK; 
         if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
      }
      else if(charge==0){
         // neutral pion
         if(qtyp==UQUARK) fqtyp = UQUARK; 
         if(qtyp==DQUARK) fqtyp = UQUARK;
         if(qtyp==UBAR) fqtyp = UQUARK;
         if(qtyp==DBAR) fqtyp = UQUARK; 
         
         if((qtyp==SBAR)||(qtyp==SQUARK))
            fqtyp = DQUARK; // always unfavored
      }
      
      // Always called for PIPLUS, either UQUARK (favored)
      // or DQUARK (unfavored)
      int icp = 1; 
      fdss_(&ih,&icp,&iset,&z,&qs2,&fqtyp,&dh);
      
      // Protection from negative values...
      
      if(dh < 0.0) dh = 0.0; 
      
      const VALUES& params = gParams->Collins.parameters;
      
      // Need to calculate and factor in the Collins part
      // of the fragmentation function
      double ptq2_average = gParams->GetPtq2Average();
      double collins_width, integrand; 
      
      switch(fqtyp) {

         case(UQUARK): // FAVORED:
            if(atPosivitityBound){
               collins = 2.*dh;
            }
            else {
               
               collins = 2.*dh*collins_z_dependence(z, params.a_up, 
                                                    params.b_up, 
                                                    params.n_up);
               //               collins_width = params.m2_up*ptq2_average/(ptq2_average + params.m2_up);
               // This is only z-dependent Collins, so integrate out pT
               collins_width = params.m2_up * ptq2_average/(ptq2_average + params.m2_up);
               integrand = sqrt( TMath::Pi() * collins_width )/2./2.; // Mulders' definition
               integrand *=  4.; // See Eq (45.) Delta D = 4 * H(1/2) 
               
               collins *= integrand; 
            } // else
            break;
            
         case(DQUARK): // UNFAVORED:
            if(atPosivitityBound){
               collins = -2.*dh; 
            }
            else {
               collins = 2.*dh*collins_z_dependence(z, params.a_down, 
                                                    params.b_down, 
                                                    params.n_down);
               //               collins_width = params.m2_down*ptq2_average/(ptq2_average + params.m2_down);
               // This is only z-dependent Collins, so integrate out pT
               collins_width = params.m2_down * ptq2_average / (ptq2_average + params.m2_down);
               integrand = sqrt( TMath::Pi() * collins_width )/2./2.; // Mulders' definition
               integrand *=  4.; // See Eq (45.) Delta D = 4 * H(1/2) 
               
               collins *= integrand; 
            }
            break;
            
         default:
            collins = 0.0; 
            break;
      }
      // Why is there no "spin" factor (sign) in this equation? Because what we want is the angular 
      // distribution referenced to the spin direction. This is accounted for properly 
      // when we calculate the angular distrubution of the hadrons in the fragmentation
      // because it is done with respect to the spin direction.
      
      // NOTE ALSO: The angular distribution that is given is with respect to the vector
      // p^ x jT^, so the dot product gives a cos(phi) behavior. Because this is perp. to
      // jt^ the hadron distribution will have a sin(phi) modulation.
      double full_FF = dh + (0.5*cos(phi)*depol_factor*collins); 
      return std::max(full_FF, 0.);
   }
}

namespace tppmc {
   
   double CollinsFragmentation::Evaluate(const pythpart& parton,
                                         const Hadron& hadron,
                                     double scale) {
      if(not Supports(parton)) {
         return mDss.Evaluate(parton, hadron, scale);
      } // if
      
      switch(hadron.Type()) {
         case  211:
         case -211: // Deliberate fall-through
         case  111: // Deliberate fall-through
            break;
         default:
            return mDss.Evaluate(parton, hadron, scale);
            break;
      } // switch(h)

      Function()->SetParameter(SCALE,  sqrt(scale));
      Function()->SetParameter(PARTON, kkp::fromPdg(parton.type));
      mFunction->SetParameter(6, parton.cos_spin_depol);
      
      SetFunctionParameters(hadron.type);
      return Function()->Eval(hadron.z, hadron.PxPyPz().Phi());
   }
   
   CollinsFragmentation::CollinsFragmentation(const FragmentationFunction* ff,
                                              TVector3 spin,
                                              int order)
   : mUnpolarised(ff)
   , mSpin(spin.Unit()) // Normalise the spin vector
   , mAtPositivityBound(false)
   , mDss(order, 0.2)
   , mFunction(NULL) {
      mParameters.Collins.SetReadFile("tables/collins_parameters.dat");
      mParameters.Collins.ReadParameters();
      gParams = &mParameters;
   }

   CollinsFragmentation::~CollinsFragmentation() {
      if(mFunction) {
         delete mFunction;
         mFunction = NULL;
      } // if
   }

   TF2* CollinsFragmentation::Function() {
      if(not mFunction) {
         mFunction = new TF2J("tppmc_collins_frag_function",
                              DSSFragFunctionCollins,
                              0.05, 1.,
                              -TMath::Pi(), TMath::Pi(),
                              8);
         // Set default parameters:
         mFunction->SetParameter(0, 1); // This isn't actually used any more!
         mFunction->SetParameter(1, 2.5); // Scale Q (not Q^2)
         mFunction->SetParameter(2, UQUARK); // Parton type
         mFunction->SetParameter(3, LO); // Calculation order
         mFunction->SetParameter(4, 1); // Hadron type
         mFunction->SetParameter(5, 1); // Charge
         mFunction->SetParameter(6, 1.); // Depolarisation factor
         mFunction->SetParameter(7, mAtPositivityBound);
         // Set number of points in x, y so TF2 interpolation is better
         mFunction->SetNpx(100);
         mFunction->SetNpy(100);
      } // if
      return mFunction;
   }

   void CollinsFragmentation::Qa(std::string filename) {
      TCanvas canvas;
      canvas.Print((filename + "[").c_str());
      int partons[] = { 2, 1 };
      int hadrons[] = { 211, -211, 111 };
      float depols[] = { 0.1, 0.3, 0.5, 0.7 };
      std::ostringstream oss;
      BOOST_FOREACH(int parton, partons) {
         BOOST_FOREACH(int hadron, hadrons) {
            BOOST_FOREACH(double depol, depols) {
               SetFunctionParameters(hadron);
               Function()->SetParameter(1, 10.);
               Function()->SetParameter(2, kkp::fromPdg(parton));
               Function()->SetParameter(6, depol);
               oss << "Parton " << parton << " hadron " << hadron << " depol " << depol;
               TH1D hist("", oss.str().c_str(), 9, 0.05, 0.95);
               oss.str("");
               for(int i(1); i <= 9; ++i) {
                  double integral = Function()->Integral(0.05, 0.05 + i * 0.1);
                  hist.SetBinContent(i, integral);
               } // for
               hist.Scale(1. / hist.GetBinContent(hist.GetNbinsX()));
               hist.SetMinimum(0.);
               hist.Draw();
               canvas.Print(filename.c_str());
            }
         }
      }
      canvas.Print((filename + "]").c_str());
   }

   std::vector<Hadron>
   CollinsFragmentation::Fragment(const pythpart& parton, double scale) {
      
      if(not Supports(parton)) {
         return mDss.Fragment(parton, scale);
      } // if
      double pt = std::max(scale, sqrt(MinimumScale()));
      
      Function()->SetParameter(1, pt);
      Function()->SetParameter(2, kkp::fromPdg(parton.type));
      
      int h = ThrowHadronType(parton, scale);
      switch(h) {
         case  211:
         case -211: // Deliberate fall-through
         case  111: // Deliberate fall-through
            break;
         default:
            return mDss.Fragment(parton, scale);
            break;
      } // switch(h)
      
      Hadron hadron;
      hadron.type = h;
      
      SetFunctionParameters(h);
      
      // Find the proton beam spin in the parton coordinate frame.
      // Set function parameter 6 to the spin component transverse
      // to the parton direction.
      if(parton.cos_spin_depol > 0. and parton.cos_spin_depol <=1.) {
         Function()->SetParameter(6, parton.cos_spin_depol);
      }
      else {
         mFunction->SetParameter(6, 0.);
      }
      double z(-1.), phi(-1.);
      while(z > mZUpper or z < mZLower) {
         Function()->GetRandom2(z, phi);
      } // while
      hadron.z = z;
      hadron.phigen = phi;
      hadron.mParton = parton;
      
      static const double MEAN_JT2 = 0.2;
      double jt_sigma = sqrt(MEAN_JT2/2.0);  
      double jt; 
      // Throw jt according to an exponential distribution such
      // that mean jt^2 is 0.2
      do {
         jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
      } while(jt>(3.0*jt_sigma));
      
      // Randomly generate phi.
      TVector3 mom(jt, 0., 0.);
      mom.SetPhi(phi);
      // TVector3::SetPhi rotates phi leaving magnitude and theta unchanged.
      hadron.px = mom.Px();
      hadron.py = mom.Py();
      
      if(211 == h and UQUARK == parton.Pdg() and 1 == parton.Spin()) {
         cffPiU()->Fill(z, phi);
      } // if
      if(111 == h and DBAR == parton.Pdg() and -1 == parton.Spin()) {
         cffPi0DbarSd()->Fill(z, phi);
      } // if
      
      return std::vector<Hadron>(1, hadron);
   }

   TH2* CollinsFragmentation::cffPiU() {
      
      static TH2D* cff_pi_u(NULL);
      if(not cff_pi_u) {
         cff_pi_u = new TH2D("cff_pi_u",
                             "Collins FF as sampled (pi+ from u quark, spin up)",
                             100,0.0, 1.0,100,0.0,TMath::TwoPi());
      } // if
      return cff_pi_u;
   }

   TH2* CollinsFragmentation::cffPi0DbarSd() {
      
      static TH2D* cff_pi0_dbar_sd(NULL);
      if(not cff_pi0_dbar_sd) {
         cff_pi0_dbar_sd = new TH2D("cff_pi0_dbar_sd",
                                    "Collins FF as sampled (pi0 from dbar quark, spin down)",
                                    100,0.0, 1.0,100,0.0,TMath::TwoPi());
      } // if
      return cff_pi0_dbar_sd;
   }

   TVector3
   CollinsFragmentation::RotateToPartonFrame(const pythpart& parton,
                                             const TVector3& v) const {
      // Generate the rotation matrix that will put us into the 
      // fragmenting parton frame:
      
      TRotation rotM;
      rotM.SetZAxis(parton.PxPyPz());
      return rotM.Inverse() * v;
   }

   bool CollinsFragmentation::Supports(const pythpart& parton) const {
      return
      (abs(parton.Pdg()) < 4 and parton.Pdg() not_eq 0);
   }

   const std::vector<int>& CollinsFragmentation::Hadrons() const {
      static std::vector<int> hadrons;
      if(hadrons.empty()) {
         int array[3] = {
            211, -211, 111
         };
         hadrons.assign(array, array + 3);
      }
      return hadrons;
   }
   
   std::string CollinsFragmentation::Name() const {
      std::string name("DSS + Collins");
      if(mAtPositivityBound) {
         name.append( " (at positivity bound)");
      } // if
      return name;
   }
   
} // namespace tppmc
